from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage, AIMessage

@tool
def add(a: int, b: int) -> int:    
    """Adds two integers and returns the result."""    
    return a + b

@tool
def multiply(a: int, b: int) -> int:    
    """Multiplies two integers and returns the result."""    
    return a * b

# case 1: run tools directly with .invoke() method
print("============= case 1: run tools directly  ==============")
result = multiply.invoke(input={"a": 6, "b": 7})  # returns 42
print(result)  # prints 42

# case 2: run tools directly with a ToolMessage object
print("============= case 2: run tools with tool message  ==============")
tool_call = {
    "type": "tool_call",
    "id": "1a",
    "args": {"a": 9, "b": 7}
}
result = multiply.invoke(input=tool_call) # returns a ToolMessage object
print(result)  # prints 42

# run tools in an agent loop with a model that can call tools
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent
from langchain_ollama import ChatOllama

llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

# run tools with prebuild agent
agent = create_react_agent(
    model=llm,
    tools=[add, multiply]
)
print("============= case 3: run tools in an agent loop  ==============")
result = agent.invoke({"messages": [
    # {"role": "user", "content": "what's 3 x 7?"},
    HumanMessage(content="what's 3 + 5?")
]})
for msg in result['messages']:
    # print(f" {getattr(msg, 'content', '')} with tool_call={getattr(msg, 'tool_calls', None)}")
    print(f"{msg.__class__.__name__} : {msg}")

# model_with_tools = llm.bind_tools([multiply])

# response_message = model_with_tools.invoke("what's 42 x 7?")
# tool_call = response_message.tool_calls[0]
# print(tool_call)  # prints ToolMessage object with tool call details
# print("=============  run tools in an agent loop  ==============")
# print(response_message.tool_calls)  # prints ToolMessage object with tool call details
# tool_msg = multiply.invoke(tool_call)
# print(f"{tool_msg.__class__.__name__} : {tool_msg}")  # prints 294
